from pymilvus import Collection
from ..milvus.module import Milvus



class MilvusHybrid(Milvus):
    def __init__(self, metric, dim, index_param):
        super().__init__(metric, dim, index_param)
        self._index_m = index_param.get("M", None)
        self._index_ef = index_param.get("efConstruction", None)
        self._selectivity_min = 0
        self._selectivity_max = 0
        self._data_count = 0

    def get_index_param(self):
        return {
            "index_type": "HNSW",
            "metric_type": self._metric_type,
            "params": {
                "M": self._index_m,
                "efConstruction": self._index_ef
            }
        }

    def set_query_arguments(self, ef, selectivity):
        """设置查询参数，包括ef和选择性范围，以及过滤模式

        Args:
            ef: HNSW算法的ef搜索参数
            selectivity: 选择范围，默认0-100
        """
        # 转换为0-1范围的小数
        self._selectivity_min = selectivity[0] / 100
        self._selectivity_max = selectivity[1] / 100

        # 计算数据总量（如果尚未计算）
        if self._data_count == 0:
            self._data_count = self.get_row_count()

        # 设置搜索参数
        self.search_params = {
            "metric_type": self._metric_type,
            "params": {"ef": ef}
        }

    def query(self, v, n):
        """使用选择性范围限制查询结果，支持标准过滤和迭代过滤

        Args:
            v: 查询向量
            n: 需要返回的最近邻数量

        Returns:
            list: 最近邻ID列表
        """
        # 确保collection已初始化
        if not self.collection:
            try:
                if utility.has_collection(self.collection_name):
                    self.collection = Collection(
                        self.collection_name,
                        consistence_level="STRONG"
                    )
                    print(f"[Milvus] 初始化集合 {self.collection_name} 成功")
                else:
                    print(f"[Milvus] 集合 {self.collection_name} 不存在，无法查询")
                    return []
            except Exception as e:
                print(f"[Milvus] 初始化集合时出错: {e}")
                return []
        # 计算ID范围限制
        min_id = int(self._data_count * self._selectivity_min)
        max_id = int(self._data_count * self._selectivity_max)

        # 搜索参数，包括基本参数和附加提示
        search_params = self.search_params.copy() if self.search_params else {"metric_type": self._metric_type}

        try:
            results = self.collection.search(
                param=search_params,
                anns_field="vector",
                data=[v],
                expr=f"id >= {min_id} and id <= {max_id}",  # 使用filter而不是expr
                limit=n,
                output_fields=["id"]
            )

            ids = [r.get("id") for r in results[0]]
            return ids
        except Exception as e:
            print(f"[Milvus] 查询时出错: {e}")
            return []
        

    def __str__(self):
        """返回对象的字符串表示"""
        return f"milvus-hybrid(ef={self.search_params.get('params', {}).get('ef') if self.search_params else 'None'}, selectivity_range={int(self._selectivity_min * 100)}%-{int(self._selectivity_max * 100)}%)"


